/*
 * Molecule.cpp
 *
 *  Created on: Feb 8, 2011
 *      Author: shaide
 */

#include "Molecule.h"
using namespace std;


#define NDEBUG

/**
 * Construct a Molecule object from an obmol
 */
Molecule::Molecule(OpenBabel::OBMol &obmol) : _size(obmol.NumAtoms())
{
	using namespace std;

	OpenBabel::OBAtomIterator it = obmol.BeginAtoms();

	//Initialize the adjecancy matrix
	for (int i=0 ; i<_size ; i++)
	{
		_adjecancy.push_back(new AtomSet(_size));
	}

	//Push all atoms in
	do
	{
		_atoms.push_back(new Atom(*it)); //populate the atom vector
		for (OpenBabel::OBBondIterator itr = (*it)->BeginBonds(); itr != (*it)->EndBonds(); itr++)
		{	//populate the adjacency matrix
			_adjecancy[(*itr)->GetEndAtomIdx()-1]->set((*itr)->GetBeginAtomIdx()-1);
			_adjecancy[(*itr)->GetBeginAtomIdx()-1]->set((*itr)->GetEndAtomIdx()-1);
		}
		assert(_size>=0);
	} while ((*it++)->GetIdx() < (unsigned int) _size);

	//Call init similarity
	cout << "Calling initSimilarity..." << endl;
	_classes.resize(initSimilarity(DEPTH_ITERATIONS));
	cout << "We have " << numClasses() << " equivalence classes" << endl;

	//Use class labels to populate the SymmetryClass structs
	for (int i=0; i<_size; i++)
	{
		_classes[_atoms[i]->_class]._atoms.push_back(_atoms[i]);
		if(_classes[_atoms[i]->_class]._set == NULL) _classes[_atoms[i]->_class]._set = new AtomSet(_size);
		_classes[_atoms[i]->_class]._set->flip(i);
	}

	//Normalize
	if (normalize()){
		//TODO Exception??
		//		if (writeOpenu) { //TODO is writeOpenU is a flag???
		//			printf("ERR* Failed to normalize atom positions: dimension of set of points = zero *ERR\n");
		//		} else {
		cerr << "Failed to normalize atom positions: dimension of set of points = zero" << endl;
		//		}
		exit(1);//TODO not good need exception or return value!!
	}
}

void Molecule::printAtoms(AtomSet* set, int from, int to)
{
	using namespace std;
	cout << "Size = " << _size << "\tClasses: " << numClasses() << "\tMax class:" << getMaxClassSize() << endl;
	assert(to <= _size);
	for (int i = from ; i < ((to==-1) ? _size : to) ; i++) {
		if(set == NULL || (*set)[i])
		{
			cout << i << ": ";
			_atoms[i]->printAtom();
		}
	}
}

Molecule::~Molecule()
{
	_atoms.clear();
}


//Returns a bitmask representation of the set
//If property = null or no function is given will return an empty set
//If property != null will perform action on any matching atom as long as the atom isn't null
//The match corresponds with the state of the atoms before any change
AtomSet Molecule::getSetMask(std::function<bool (const Atom)> property, std::function<void (Atom&)> action = NULL)
{
	AtomSet ret(_size);
	if (property != NULL) for(int i=0; i<_size ; i++)
	{
		bool match = property(*_atoms[i]);
		if (match)
		{
			ret.set(i, true);
			if (action != NULL) action(*_atoms[i]);
		}
	}
	return ret;
}

//Finds any atom which satisfies property and returns it's index
//If action != null it will be performed on the atom
//The atom index will be returned even if it doesn't match the property after the action was applied
int Molecule::findAtom(std::function<bool (const Atom)> property, std::function<void (Atom&)> action = NULL)
{
	for (int i=0 ; i<_size ; i++)
	{
		if (property(*_atoms[i]))
		{
			if (action != NULL) action(*_atoms[i]);
			return i;
		}
	}
	return -1;
}

bool Molecule::areAdj(int A, int B)
{
	bool val = _adjecancy[A]->test(B);
	return (val);
}

bool Molecule::isSimilar(const Atom &A, const Atom &B)
{
	if (strcmp(A._symbol,B._symbol) != 0 ) return false;
	if (A._valence != B._valence) return false;
	AtomSet S(_size);
	for (int i = 0; i < _size; i++)
	{
		if (!areAdj(A._index,i)) continue; //Skip any atoms which aren't adjecant to A
		int K = findAtom([&](const Atom& b)
				{
			return areAdj(B._index,b._index) && (_atoms[i]->_class == b._class) && (!S[b._index]);
				});
		if (K == -1) return false;
		S.flip(K);
	}
	return true;
}

bool Molecule::normalize()
{

	double x_avg, y_avg, z_avg,norm;
	x_avg = y_avg = z_avg = 0.0;

	double mass_sum = 0;

	std::vector<Atom*>::iterator it = _atoms.begin();
	while (it != _atoms.end())
	{
		x_avg += (*it)-> _pos[0] * (*it)-> _mass;
		y_avg += (*it)-> _pos[1] * (*it)-> _mass;
		z_avg += (*it)-> _pos[2] * (*it)-> _mass;
		mass_sum += (*it)->_mass;
		it++;
	}
	x_avg /= (double)(mass_sum);
	y_avg /= (double)(mass_sum);
	z_avg /= (double)(mass_sum);

	norm = 0.0;

	it = _atoms.begin();

	while (it != _atoms.end()){
		norm += sqrt(fabs((*it)->_pos[0] - x_avg)) + sqrt(fabs((*it)->_pos[1] - y_avg)) + sqrt(fabs((*it)->_pos[2] - z_avg));
		it++;
	}

	norm = sqrt(norm);

	if(norm < MINDOOUBLE){
		return false;
		//TODO Exception VS bool??
		//		if (writeOpenu) { //TODO is writeOpenU is a flag???
		//			printf("ERR* Failed to normalize atom positions: dimension of set of points = zero *ERR\n");
		//		} else {
		//		cerr << "Failed to normalize atom positions: dimension of set of points = zero" << endl;
		//		}
		//		exit(1);//TODO not good need exception or return value!!
	}

	it = _atoms.begin();

	while (it != _atoms.end()){
		(*it)->_pos[0] = ((*it)->_pos[0] - x_avg) / norm;
		(*it)->_pos[1] = ((*it)->_pos[1] - y_avg) / norm;
		(*it)->_pos[2] = ((*it)->_pos[2] - z_avg) / norm;
		it++;
	}
	_norm = norm;

	return true;

}

int Molecule::initSimilarity(int depth)
{
	cout << "_size is: " << _size << "\t depth is: " << depth << endl; //TODO delete
	AtomSet viewing(_size);
	int m = 0;
	{
		// break into initial groups by symbol and valancy
		for(int i=0 ; i<_size ; i++)
		{
			if(viewing[i]) continue;
			AtomSet mask = getSetMask(
					[&](const Atom& b)->bool
					{
				return (!viewing[b._index] && b._valence == _atoms[i]->_valence && strcmp(b._symbol, _atoms[i]->_symbol) == 0);
					},

					[&](Atom &b)
					{
						b._class = m; //set the equivalence class of atom 'b' to 'm'
					}
			);
			m++;
			viewing |= mask; //viewing now contains a set of atoms and for each atom it's equivalence class
		}
		cout << "number of equivalence class: " << m << endl; //TODO DELETE!
	}

	// iteratively refine the breakdown into groups
	for (int j = 0; j < depth ; j++)
	{
		AtomSet viewed(_size);

		int before = m;
		// iteratively refine the breakdown into groups
		for (int i = 0 ; i < _size ; i++) 	//for any atom A
		{
			if (viewed[i]) continue;
			viewed.flip(i);
			AtomSet mask = getSetMask(//put all atoms which are equivalnt but not similar to A in
					//their own equivalence class
					[&](const Atom& b)->bool
					{
				if (viewed[b._index]) return false; //if b is in viewed return false;

				if (_atoms[i]->_class == b._class) //if in the same class add b to viewed
				{
					viewed.flip(b._index);
					bool similar = !isSimilar(*_atoms[i],b);
					return similar;
				}
				return false;
					},
					[&m](Atom& b)
					{
						b._class = m;
					}
			);
			if (!mask.none()) m++;
		}
		if (before == m){
			std::cout << "Finished early after just " << j << " iterations" << std::endl;
			return m;
		}
	}
	return m;
}

Molecule::SymmetryClass& Molecule::getClass(int classId)
{
	return _classes[classId];
}

int Molecule::getClassSize(int classId)
{
	return _classes[classId]._atoms.size();
}

int Molecule::numClasses()
{
	return _classes.size();
}


int Molecule::getMaxClassSize()
{
	int size=0;
	unsigned int i=0;
	for (i=0 ; i<_classes.size() ; i++)
	{
		if (getClassSize(i) > size) size = getClassSize(i);
	}
	return size;
}
